import numpy as np
import matplotlib.pyplot as plt
from joblib import dump, load
from sklearn.ensemble import RandomForestRegressor
import time
import pandas as pd
from sklearn.model_selection import train_test_split
from copy import deepcopy
import torch
# device
device = 'cuda:2'


# 构建目标函数
temperature0 = [175, 175, 175, 175, 175, 195, 235, 255, 255, 25, 25]
temperature1 = [173, 173, 173, 173, 173, 198, 230, 257, 257, 25, 25]
deltat = 0.5 # seconds
deltax = 0.0001 # milimeters
timestamps = [[0, 20], [20, 70], [70, 169.0], [169,  200.0], [200, 230.0], [230, 294], [294, 375], [375, 500]]# 最后一个时间区间不用
H = 0.015 # milimeters

def Tw(x, temperature = temperature0):
    wenqu = 30.5
    jiange = 5
    s = [25]
    for i in range(11):
        s.append(s[-1] + wenqu)
        s.append(s[-1] + jiange)
    s = s[:-1]
    #print(s)
    for i in range(0, 22, 2):
        if x < 25:
            return x * (temperature[0] - 25) / 25 + 25
        if x >= s[i] and x <= s[i + 1]:
            return temperature[int(i/2)]
        elif x > s[i + 1] and i < 20 and x < s[i + 2]:
            return (x - s[i + 1]) / jiange * (temperature[int(i/2 + 1)] - temperature[int(i/2)]) + temperature[int(i/2)]
        
    return temperature[-1]
v = 70 / 60
def model():

    df = pd.read_excel("ants_tryout/附件.xlsx")
    x = df['时间(s)']
    y = df['温度(ºC)']
    x_data = np.array(x)
    y_data = np.array(y)
    train, test, train_labels, test_labels = train_test_split(x_data.reshape(-1,1), y_data, test_size=0.2, random_state=42)
    rf = RandomForestRegressor(n_estimators=50, random_state=42)
    rf.fit(train, train_labels)
    #y_pred = rf.predict(test)
    score = rf.score(test, test_labels)
    print(f"R^2 Score: {score}")
    return rf

def loss(alphas, hs, Td, flag = 0):
    #alpha = 4.03 * 10**(-5) 
    arrt = np.arange(0, 375, deltat)
    arrx = np.arange(0, H, deltax)
    # U = np.zeros((len(arrx), len(arrt)))
    # U[:,0] = 25 # u(x_i, 0) = 25
    m = len(arrx)
    u0 = np.ones(m)
    u0 = u0 * 25
    A = np.zeros((m, m))
    A[0,0] = 0
    A[m - 1, 0] = 0
    B = np.zeros((m, m))           
    res = [25]
    res0 = [25]
    target = []
    step = 80
    B = torch.from_numpy(B).to(device)
    A = torch.from_numpy(A).to(device)
    
    # 根据alpha和h更新A, invertB
    def refresh(alpha, h):
        nonlocal A, B

        r = alpha * deltat / deltax**2
        m = B.shape[0]

        B[0,0] = 1 + h * deltax
        B[0,1] = -1
        B[m - 1, m - 1] = 1 + h * deltax
        B[m - 1, m - 2] = -1
        for i in range(1, m-1):
            A[i, i - 1] = r
            A[i, i] = 2 * (1 - r)
            A[i, i + 1] = r
        for i in range(1, m - 1):
            B[i, i - 1] = 0 - r
            B[i, i] = 2 * (1 + r)
            B[i, i + 1] = 0 - r
        B.to(device)
        A.to(device) 
        #invertB = np.linalg.inv(B)
        invertB = torch.inverse(B)
        return A, invertB
    #进行一次迭代
    def iter(u0, invertB, A, h, i):        
        def countBias(i):
            bias = np.zeros(A.shape[0])
            bias[0] = h * Tw(v * i) * deltax
            bias[m - 1] = h * Tw(v * i) * deltax
            bias = torch.from_numpy(bias).to(device)

            return bias

        u0 = torch.from_numpy(u0).to(device)
        bias = countBias(i)
        return invertB @ (A @ u0 + bias)
    
    in_flag = 0
    # alpha = alphas[0]
    # h = hs[0]
    # A, invertB = refresh(alpha, h)
    for i , curr_time in enumerate(arrt):
        if i == len(arrt) - 1:
            break
        if curr_time >= timestamps[flag][1]:
            break
        if curr_time >= timestamps[in_flag][0]:
            #print(f"refresh:{in_flag + 1}")
            if in_flag > flag:
                break
            alpha = alphas[in_flag]
            h = hs[in_flag]
            A, invertB = refresh(alpha, h)
            in_flag = in_flag + 1
            #print(arrt[i])
        iter_result = iter(u0, invertB, A, h, arrt[i + 1])
        iter_result = iter_result.cpu().numpy()
        
        # if int(len(arrx) / 2) > boundaries[flag]:
        #     refresh(alpha[flag], alpha[flag])
        #print(result)
        
        res.append(u0[int(len(arrx) / 2)])
        res0.append(u0[0])
        u0 = iter_result
        
        if arrt[i + 1] >= 19 :
            target.append(Td(arrt[i + 1]))
            # if in_flag == flag + 1:
            #     if i % step == 0:
            #         print("exp: " + str(res[-1]))
            #         print("exp_0: " + str(res0[-1]))
            #     if i % step == 0 :
            #         print("standard:" + str(target[-1]))

    res = res[-len(target):]    
    return np.sum(np.abs(np.array(res) - np.array(target)) ** 2)
    #return np.mean(np.abs(np.array(res) - np.array(target)))

rf = model()
def Td(t):
    return rf.predict([[t]])[0]

# 目标函数
def fitnessfunction(x, flag = 6):
    a = x[0]
    h = x[1]
    # [3.58636647e-07 6.08669053e-07]
    # [6.37808028e+02 7.52005205e+02]
    # a[flag - 1] =  alphas[flag - 1] 
    # h[flag - 1] =  hs[flag - 1]
    y = loss(a, h, Td, flag)

    return y,x




# 粒子中心
standard = np.array([[5.07374444e-07, 5.99417655e-07, 6.46113456e-07, 6.64988285e-07,
  9.96610937e-07, 6.90358826e-07, 3.32748743e-07],
 [8.06859815e+02, 7.48398526e+02, 6.72349848e+02, 1.02191905e+03,
  7.38533936e+02, 7.60863780e+02, 6.54877678e+02]])
# initial 粒子初态
size = 20 # 粒子数量
paticles = []
snum = 7 # 参数个数
for i in range(size):
    #paticle = np.random.rand(2,1) * np.array([[0.2e-7],[70]]) + np.array([[6.03357225e-07],[4.3351940e+01]])
    
    paticle = np.random.rand(2, snum) * np.array([[3e-7] * snum,[200] * snum]) + standard
    paticles.append(paticle)
pbest_choice = deepcopy(paticles)
paticles[0] = standard
pbest = [0] * 20# 个体最优值
gbest = 1000# 群体最优值
gbest_choice = 0# 参数取值


# 初始化控制变量
c1 = 2
c2 = 2
w_init = 0.9 
w_end = 0.4
vmax = np.array([[1e-7] * snum, [90] * snum])
vmin = np.array([[-1e-7] * snum, [-90] * snum])
paticle_min = np.array([[1e-8] * snum, [0.5] * snum])
paticle_max = np.array([[6e-6] * snum, [5000] * snum])

# 初始化每个粒子速度
ve = []
for i in range(size):
    #velocity = np.random.randn(2,1)
    velocity = np.random.randn(2, snum) * 1e-5
    ve.append(velocity)

start_time = time.time()


# 计算每个粒子的loss值，初始化群体最优值
for i in range(size):
    pbest[i], _ = fitnessfunction(paticles[i])
    if pbest[i] < gbest:
        gbest = pbest[i]
        gbest_choice = paticles[i]


def PSO(iteration):
    global gbest,gbest_choice
    i = 0
    print("begin")
    while True:
        i += 1
        print("迭代次数：",i)
        print("最好结果：",gbest)
        print("最优取值：",gbest_choice)
        print("用时：", time.time() - start_time)
        if gbest < 0.3 or i == iteration:
            break
        w= (w_init - w_end) * (iteration - i) / iteration + w_end
        
        for j in range(size):
            #更新位置和速度，下面就是我们之前重点讲解的两条公式。
            ve[j]=w * ve[j]+c1*np.random.rand()*(pbest_choice[j]-paticles[j])+c2*np.random.rand()*(gbest_choice-paticles[j])

            # print("item1:")
            # print(c1*np.random.rand()*(pbest_choice[j]-paticles[j]))
            # print("item2:")
            # print(c2*np.random.rand()*(gbest_choice-paticles[j]))
            # print("item0")
            # print(w*ve[j])
            if (ve[j]>vmax).any():
                ve[j]  = np.minimum(ve[j],vmax)  
            
            if (ve[j] < vmin).any():
                ve[j] =  np.maximum(ve[j],vmin) 
                
            
            
            # print(paticles[j])
            paticles[j] = paticles[j] + ve[j]


            countmin = 0
            countmax = 0
            #越界判断，范围限定在[0, 2]
            while (paticles[j] > paticle_max).any():
                if countmax > 10:
                    
                    break
                countmax = countmax + 1
                print("!!!")
                #new_pixel = np.random.rand(2,1) * np.array([[0.2e-7],[1]])
                new_pixel = np.random.rand(2,snum) * np.array([[5e-8] * snum,[10] * snum])
                paticles[j] = new_pixel + pbest_choice[j] if np.random.randn() > 0 else -1 *  new_pixel + pbest_choice[j]
            while (paticles[j] < paticle_min).any():
                if countmin > 10:
                    paticles[j] = paticle_min
                    break
                countmin = countmin + 1
                print("!!!")
                #new_pixel = np.random.rand(2,1) * np.array([[0.2e-7],[1]])
                new_pixel = np.random.rand(2,snum) * np.array([[0.2e-7] * snum,[10] * snum])
                paticles[j] = new_pixel + pbest_choice[j] if np.random.randn() > 0 else -1 *  new_pixel + pbest_choice[j]
                
            
            
        #更新个体极值和群体极值
        for j in range(size):
            # print(paticles[j])
            y, choice = fitnessfunction(paticles[j])
            if pbest[j] > y:
                pbest[j] = y
                pbest_choice[j] = paticles[j]
            
            # print(paticles[j])
            # print(pbest[j])
            if pbest[j]<gbest:
                gbest=pbest[j]
                gbest_choice = pbest_choice[j]
                # print(gbest)
                # best_choice = choice
                #System.out.println("粒子n"+j+": x = "+x[j]+"  "+"v = "+v[j]);
            
            #System.out.println("第"+(i+1)+"次迭代，全局最优解 gbest = "+gbest);
            #System.out.print("\n");
        
    return gbest, gbest_choice





best, best_choice = PSO(100)
print(best)
print(best_choice)

